#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
RRT路径规划应用程序 - 主入口点

此脚本作为RRT路径规划应用程序的主入口点。
它负责解析命令行参数、加载配置、
初始化环境和规划器，并执行规划过程。
"""

import os
import argparse
import yaml
import time
import numpy as np
from datetime import datetime

from environment.map import Map
from robot.robot_model import Robot
from planning.planner_factory import create_planner
from utils.plotter import Plotter
from utils.collision_checker import CollisionChecker


def parse_arguments():
    """解析命令行参数。"""
    parser = argparse.ArgumentParser(description='RRT Path Planning Application')
    parser.add_argument(
        '--config', 
        type=str, 
        default='configs/default_config.yaml',
        help='配置文件路径'
    )
    parser.add_argument(
        '--output_dir', 
        type=str, 
        default='results',
        help='保存结果的目录'
    )
    parser.add_argument(
        '--visualize', 
        action='store_true',
        help='在规划过程中启用可视化'
    )
    parser.add_argument(
        '--save_animation', 
        action='store_true',
        help='保存规划过程的动画'
    )
    parser.add_argument(
        '--seed', 
        type=int, 
        default=None,
        help='用于可重现性的随机种子'
    )
    parser.add_argument(
        '--adaptive', 
        action='store_true',
        help='启用自适应采样（覆盖配置设置）'
    )
    
    return parser.parse_args()


def load_config(config_path):
    """从YAML文件加载配置。"""
    with open(config_path, 'r') as file:
        config = yaml.safe_load(file)
    return config


def setup_output_directories(base_dir):
    """设置结果输出目录。"""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    result_dir = os.path.join(base_dir, f"run_{timestamp}")
    
    paths_dir = os.path.join(result_dir, "paths")
    plots_dir = os.path.join(result_dir, "plots")
    stats_dir = os.path.join(result_dir, "stats")
    
    os.makedirs(paths_dir, exist_ok=True)
    os.makedirs(plots_dir, exist_ok=True)
    os.makedirs(stats_dir, exist_ok=True)
    
    return {
        "base": result_dir,
        "paths": paths_dir,
        "plots": plots_dir,
        "stats": stats_dir
    }


def main():
    """运行RRT路径规划应用程序的主函数。"""
    # 解析命令行参数
    args = parse_arguments()
    
    # 加载配置
    config = load_config(args.config)
    
    # 如果提供了随机种子则设置
    if args.seed is not None:
        np.random.seed(args.seed)
    else:
        # 如果可用，从配置中设置种子
        seed = config.get("general", {}).get("seed")
        if seed is not None:
            np.random.seed(seed)
        
    # 设置输出目录
    output_dirs = setup_output_directories(args.output_dir)
    
    # 用命令行参数覆盖配置
    if args.visualize:
        if "visualization" not in config:
            config["visualization"] = {}
        config["visualization"]["enable"] = True
    
    if args.save_animation:
        if "visualization" not in config:
            config["visualization"] = {}
        config["visualization"]["save_animation"] = True
    
    # 如果通过命令行请求，启用自适应采样
    if args.adaptive:
        if "planner" not in config:
            config["planner"] = {}
        if config["planner"].get("type") == "rrt_improved":
            config["planner"]["explore_strategy"] = "adaptive"
            # 如果不存在，设置默认自适应参数
            if "adaptive" not in config["planner"]:
                config["planner"]["adaptive"] = {
                    "bias_min": 0.05,
                    "bias_max": 0.8,
                    "radius_factor": 2.0,
                    "voronoi_bias": 0.3
                }
            print("已启用自适应采样策略")
    
    # 将配置保存到结果目录
    with open(os.path.join(output_dirs["base"], "config_used.yaml"), 'w') as f:
        yaml.dump(config, f, default_flow_style=False)
    
    # 初始化环境
    map_config = config.get("environment", {}).get("map", {})
    map_instance = Map(
        width=map_config.get("width", 800),
        height=map_config.get("height", 600),
        obstacles=map_config.get("obstacles", [])
    )
    
    # 初始化机器人
    robot_config = config.get("robot", {})
    robot = Robot(
        radius=robot_config.get("radius", 10.0),
        step_size=robot_config.get("step_size", 5.0)
    )
    
    # 初始化碰撞检查器
    collision_checker = CollisionChecker(map_instance, robot)
    
    # 初始化规划器
    planner_config = config.get("planner", {})
    planner = create_planner(
        planner_type=planner_config.get("type", "rrt_base"),
        collision_checker=collision_checker,
        config=planner_config
    )
    
    # 如果启用可视化，初始化绘图器
    plotter = None
    viz_enabled = config.get("visualization", {}).get("enable", False) or args.visualize
    viz_save = config.get("visualization", {}).get("save_animation", False) or args.save_animation
    
    if viz_enabled or viz_save:
        save_dir = output_dirs["plots"] if viz_save else None
        plotter = Plotter(
            map_instance=map_instance,
            robot=robot,
            save_dir=save_dir,
            config=config  # 将完整配置传递给绘图器以进行可视化设置
        )
    
    # 提取起始和目标位置
    start = tuple(planner_config.get("start", [50, 50]))
    goal = tuple(planner_config.get("goal", [750, 550]))
    
    # 运行规划过程
    print(f"使用{planner_config.get('type', 'default')}算法规划从{start}到{goal}的路径...")
    if planner_config.get("explore_strategy") == "adaptive":
        print("使用自适应采样策略")
        
    start_time = time.time()
    
    # 对于自适应采样，将狭窄通道传递给绘图器
    narrow_passages = []
    
    # 规划路径
    path, info = planner.plan(
        start=start,
        goal=goal,
        max_iterations=planner_config.get("max_iterations", 10000),
        goal_tolerance=planner_config.get("goal_tolerance", 5.0),
        plotter=plotter
    )
    
    # 如果使用自适应采样，更新狭窄通道以进行可视化
    if hasattr(planner, 'narrow_passage_regions'):
        narrow_passages = planner.narrow_passage_regions
    
    planning_time = time.time() - start_time
    print(f"规划在{planning_time:.2f}秒内完成。")
    
    # 保存结果
    if path:
        print(f"找到包含{len(path)}个路径点的路径")
        
        # 将路径保存到文件
        import json
        path_file = os.path.join(output_dirs["paths"], "path.json")
        with open(path_file, 'w') as f:
            json.dump({"path": [list(point) for point in path]}, f, indent=2)
        
        # 保存统计数据
        stats = {
            "planning_time": planning_time,
            "path_length": planner.compute_path_length(path) if hasattr(planner, 'compute_path_length') else sum(
                np.linalg.norm(np.array(path[i]) - np.array(path[i-1])) for i in range(1, len(path))
            ),
            "path_points": len(path),
            "tree_nodes": info.get("num_nodes", 0),
            "iterations": info.get("iterations", 0),
            "narrow_passages_detected": len(narrow_passages) if narrow_passages else 0
        }
        
        if hasattr(planner, 'current_goal_bias'):
            stats["final_goal_bias"] = planner.current_goal_bias
            
        if hasattr(planner, 'rewire_radius'):
            stats["final_rewire_radius"] = planner.rewire_radius
        
        stats_file = os.path.join(output_dirs["stats"], "stats.json")
        with open(stats_file, 'w') as f:
            json.dump(stats, f, indent=2)
        
        # 生成最终图表
        if plotter:
            plotter.plot_result(
                path, 
                info.get("nodes", []),
                narrow_passages=narrow_passages
            )
            plotter.save_figure(os.path.join(output_dirs["plots"], "final_path.png"))
            
            if viz_enabled:
                print("显示最终路径图。关闭窗口以退出。")
                plotter.show()
            
            # 最终图表更新以确保目标可见
            plotter.plot_result(path, info.get("nodes", []))
            print("规划完成。关闭可视化窗口以退出。")
            plotter.show()  # 这将保持窗口打开直到用户关闭它
            
    else:
        print("未找到路径！")

    print(f"结果已保存到 {output_dirs['base']}")


if __name__ == "__main__":
    main() 